////
////  UICollectionView+Extension.swift
////  Skeleton
////
////  Created by bin hu on 2025/6/4.
////
//

import UIKit

private var skeletonDataSourceKey: UInt8 = 0
private var originalDataSourceKey: UInt8 = 1

extension UICollectionView {
    private var skeletonDataSource: SkeletonCollectionViewDataSource? {
        get { objc_getAssociatedObject(self, &skeletonDataSourceKey) as? SkeletonCollectionViewDataSource }
        set { objc_setAssociatedObject(self, &skeletonDataSourceKey, newValue, .OBJC_ASSOCIATION_RETAIN) }
    }
    
    private var originalDataSource: UICollectionViewDataSource? {
        get { objc_getAssociatedObject(self, &originalDataSourceKey) as? UICollectionViewDataSource }
        set { objc_setAssociatedObject(self, &originalDataSourceKey, newValue, .OBJC_ASSOCIATION_RETAIN) }
    }
    
    func showSkeletonWithDataSource() {
        guard let dataSource = self.dataSource else { return }
        
        // 保存原始数据源
        originalDataSource = dataSource
        skeletonDataSource = SkeletonCollectionViewDataSource(original: dataSource)
        
        // 切换为骨架数据源
        self.dataSource = skeletonDataSource
        DispatchQueue.main.async { self.reloadData() }
        
        self.isScrollEnabled = false
    }
    
    func hideTabSkeleton() {
        guard skeletonDataSource != nil else { return }
        
        // 先更新所有可见cell的状态
        visibleCells.forEach { $0.contentView.skeletonEnable = false }
        
        // 恢复原始数据源
        self.dataSource = originalDataSource
        originalDataSource = nil
        skeletonDataSource = nil
        
        DispatchQueue.main.async {
            self.reloadData()
            // 确保新加载的cell也不会显示骨架
            self.visibleCells.forEach { $0.contentView.skeletonEnable = false }
        }
        self.isScrollEnabled = true
    }
    
}

private class SkeletonCollectionViewDataSource: NSObject, UICollectionViewDataSource {
    weak var originalDataSource: UICollectionViewDataSource?
    
    init(original: UICollectionViewDataSource) {
        self.originalDataSource = original
        super.init()
    }
    
    func numberOfSections(in collectionView: UICollectionView) -> Int {
        // 保持与原数据源相同的section数
        return originalDataSource?.numberOfSections?(in: collectionView) ?? 1
    }
    
    func collectionView(_ collectionView: UICollectionView, numberOfItemsInSection section: Int) -> Int {
        // 保持与原数据源相同的行数
        return originalDataSource?.collectionView(collectionView, numberOfItemsInSection: section) ?? defaultNumbers
    }
    
    func collectionView(_ collectionView: UICollectionView, cellForItemAt indexPath: IndexPath) -> UICollectionViewCell {
        // 获取原始cell但清除内容
        guard let cell = originalDataSource?.collectionView(collectionView, cellForItemAt: indexPath) else {
            return UICollectionViewCell()
        }
        cell.contentView.showSkeleton()
        return cell
    }

}
